import pandas as pd
import numpy as np
from interpret.glassbox import ExplainableBoostingRegressor
from sklearn.preprocessing import MinMaxScaler
from interpret import show
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
from sklearn.model_selection import train_test_split

# ========== 1. 数据读取 ==========
df = pd.read_csv(r"D:\resources\data\jumbo\2014-2020_1deg_st_merged_st_chl_ssh_mld.csv")  # 包含 CPUE 和海温特征
target_col = "CPUE"

# ========== 2. 构建两种特征方案 ==========
feature_set_1 = ['Year', 'Month', 'Lon', 'Lat',
                 'ST_0.5', 'ST_47.4', 'ST_92.3', 'ST_155.9', 'ST_222.5', 'ST_318.1', 'ST_453.9']  # 垂直海温平均值
# feature_set_1 = ['Year', 'Month', 'Lon', 'Lat', 'ST_0.5', 'ST_47.4', 'ST_155.9' ]
label_cols = ['ST_0.5_Label', 'ST_47.4_Label', 'ST_92.3_Label',
              'ST_155.9_Label', 'ST_222.5_Label', 'ST_318.1_Label', 'ST_453.9_Label']
feature_set_2 = ['Year', 'Month', 'Lon', 'Lat',
                'ST_0.5_Label', 'ST_47.4_Label', 'ST_92.3_Label',
                'ST_155.9_Label', 'ST_222.5_Label', 'ST_318.1_Label', 'ST_453.9_Label']   # 多层特征
categorical_cols = [
    'ST_0.5_Label', 'ST_47.4_Label', 'ST_92.3_Label',
    'ST_155.9_Label', 'ST_222.5_Label', 'ST_318.1_Label', 'ST_453.9_Label'
]
# 将它们转换为分类变量
# for col in categorical_cols:
#     df[col] = df[col].astype('category')
# 定义特征方案
feature_sets = {
    "海温平均值方案": feature_set_1,
    "多层特征方案": feature_set_2
}

X = df[feature_set_1]
# X = pd.get_dummies(X, columns=label_cols)
y = df[target_col]
# y = df[target_col].values.reshape(-1, 1)
scaler = MinMaxScaler()
X = scaler.fit_transform(X)

# y_scaler = MinMaxScaler()
# y_scaled = y_scaler.fit_transform(y)
y = np.log1p(y)  # log(CPUE + 1)，使分布更平滑

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
# X_train, X_test, y_train, y_test = train_test_split(X, y_scaled, test_size=0.2, random_state=42)

# ========== 4. 训练 EBM 模型 ==========
# ebm = ExplainableBoostingRegressor(
#     interactions=10,        # 自动发现最多10个特征交互
#     max_leaves=5,           # 每个分段的最大叶子数
#     learning_rate=0.05,
#     min_samples_leaf=5,
#     random_state=42
# )
ebm = ExplainableBoostingRegressor(
    interactions=1,        # 自动发现最多10个特征交互
    max_leaves=5,           # 每个分段的最大叶子数
    learning_rate=0.5,
    smoothing_rounds=100,
    min_samples_leaf=5,
    # inner_bags=50,
    # early_stopping_rounds=1000,
    # max_bins=2048,
    # max_rounds=1000000000,
    random_state=42
)

ebm.fit(X_train, y_train)


# ========== 5. 预测与指标 ==========
y_pred = ebm.predict(X_test)

mse = mean_squared_error(y_test, y_pred)
rmse = np.sqrt(mse)
mae = mean_absolute_error(y_test, y_pred)
mre = np.mean(np.abs((y_test - y_pred) / (y_test + 1e-8)))  # 相对误差
r2 = r2_score(y_test, y_pred)

print("\n🎯 EBM 模型评估结果 (测试集)")
print(f"MSE  = {mse:.6f}")
print(f"RMSE = {rmse:.6f}")
print(f"MAE  = {mae:.6f}")
print(f"MRE  = {mre:.6f}")
print(f"R²   = {r2:.6f}")

# ========== 6. 可解释性分析 ==========
# --- 全局解释 ---
# ebm_global = ebm.explain_global()
# show(ebm_global)